from SiteFrequencySpectrum import *
from Parameters import *
from PolymorphismTables import PolymorphismTables
import SiteFrequencySpectrum as sfs
import infiniteSitesTools as ist

from Polyspect_Power import PowerPoint
from Polyspect_Operator import Polyspect

import numpy.matlib as np
from scipy.linalg import norm

class PolyspectCalculus():

	myLinOp = None
	myPolyspect = None
	myParameters = None
	myAltHypothesis = None
	
	myNumChromosomes = 0
	myDx = 0.01
	
	mySFSs = None
	myAltSFSs = None
	
	myPolyspectPower = None
	myLocalPowerGrid = None
	
	myDirectionalDerivatives = None
	myDirectionalConcavities = None
	
	myCalculusDone = False
	
	def __init__(self, iLinOp, iParameters, iAltHypothesis, \
				iSFS_Null = None, iSFS_Alt = None):
		self.myLinOp = iLinOp
		self.myParameters = iParameters
		self.myAltHypothesis = iAltHypothesis
		
		self.myNumChromosomes = iParameters.getNumChromosomes()
		
		if iSFS_Null is None:
			self.generateSimulationData()
		else:
			self.mySFSs = iSFS_Null
			self.myAltSFSs = iSFS_Alt
		
		self.setPolyspect( iLinOp )
	

	def setPolyspect(self, iLinOp):
		self.myPolyspect = Polyspect( iLinOp, self.myParameters, self.mySFSs )
		
		self.myPolyspectPower = PowerPoint( self.myPolyspect, self.myAltSFSs )
		print self.myPolyspectPower
		
		self.myCalculusDone = False
		
	def getKthDirectionalDerivative(self, k, whichSide="both" ):
		dx = self.myDx
		
		if self.myLocalPowerGrid[k,0] < 0.0:
			dd = np.ravel(np.zeros(self.myNumChromosomes))
			dd[k] = dx
			
			Point0 = self.myLinOp - dd
			Point1 = self.myLinOp
			Point2 = self.myLinOp + dd
	
			Point0Polyspect = Polyspect( Point0, self.myParameters, self.mySFSs )
			Point1Polyspect = Polyspect( Point1, self.myParameters, self.mySFSs )
			Point2Polyspect = Polyspect( Point2, self.myParameters, self.mySFSs )
		
			Point0Power = PowerPoint( Point0Polyspect, self.myAltSFSs )
			Point1Power = PowerPoint( Point1Polyspect, self.myAltSFSs )
			Point2Power = PowerPoint( Point2Polyspect, self.myAltSFSs )
			
			self.myLocalPowerGrid[k,0] = Point0Power.getPower()
			self.myLocalPowerGrid[k,1] = Point1Power.getPower()
			self.myLocalPowerGrid[k,2] = Point2Power.getPower()
		
		
		firstPointPower = self.myLocalPowerGrid[k,2]
		secondPointPower = self.myLocalPowerGrid[k,0]
		
		if whichSide == "less":
			secondPointPower = self.myLocalPowerGrid[k,1]
		elif whichSide == "more":
			firstPointPower = self.myLocalPowerGrid[k,1]
		else:
			dx *= 2.0
			
		deriv = (firstPointPower - secondPointPower) / dx
		
		return deriv
	
	def doCalculus(self):
		
		self.myLocalPowerGrid = np.ones((self.myNumChromosomes, 3))
		self.myLocalPowerGrid *= -1.0
		
		dd = np.zeros(( self.myNumChromosomes, 4))
		
		for k in range(1, self.myNumChromosomes):
			print "Working on " + str(k) + "-th derivative"
			dd[k, 0] = self.getKthDirectionalDerivative(k, "less")
			dd[k, 1] = self.getKthDirectionalDerivative(k, "more")
			dd[k, 2] = self.getKthDirectionalDerivative(k, "both")
			dd[k, 3] = ( dd[k,1] - dd[k,0] ) / self.myDx
			print str(dd[k,:])
			
		self.myDirectionalDerivatives = dd[:,2]
		self.myDirectionalConcavities = dd[:,3]
		
		self.myCalculusDone = True
		
	def getDerivative(self):
		if not self.myCalculusDone:
			self.doCalculus()
		
		return self.myDirectionalDerivatives
		
	def getConcavity(self):
		if not self.myCalculusDone:
			self.doCalculus()

		return self.myDirectionalConcavities

#	def generateSimulationData(self):
#		
#		file_Null = '/tmp/SFS_Null.ms'
#		file_Alt = '/tmp/SFS_Alt.ms'
#		
#		pt = PolymorphismTables()
#		msCommand = self.myParameters.getMSString_Null( self.myNumSimulations )
#		msOut = pt.runMS(msCommand, file_Null)
#		pt.readMS(file_Null, msOut)
#	
#		self.mySFSs = SiteFrequencySpectrum(pt)
#		
#		pt = PolymorphismTables()
#		msCommand = self.myParameters.getMSString_Alternative( self.myAltHypothesis, self.myNumPowerSimulation )
#		msOut = pt.runMS(msCommand, file_Alt)
#		pt.readMS(file_Alt, msOut)
#		
#		self.myAltSFSs = SiteFrequencySpectrum(pt)
		
